// src/transport_tests.rs
// Additional unit tests for transport.rs to achieve 95%+ coverage

use crate::transport::{Transport, WebSocketTransport, TcpTransport};
use tokio::net::{TcpListener, TcpStream};
use tokio_tungstenite::{connect_async, accept_async};

// Mock implementations for testing
#[cfg(test)]
mod mock {
    use crate::transport::Transport;
    use async_trait::async_trait;
    use std::io;
    
    pub struct MockTransport {
        pub read_data: Option<Vec<u8>>,
        pub write_data: Vec<u8>,
        pub is_closed_flag: bool,
    }
    
    impl MockTransport {
        pub fn new() -> Self {
            Self {
                read_data: None,
                write_data: Vec::new(),
                is_closed_flag: false,
            }
        }
        
        pub fn with_read_data(data: Vec<u8>) -> Self {
            Self {
                read_data: Some(data),
                write_data: Vec::new(),
                is_closed_flag: false,
            }
        }
        
        pub fn closed() -> Self {
            Self {
                read_data: None,
                write_data: Vec::new(),
                is_closed_flag: true,
            }
        }
    }
    
    #[async_trait]
    impl Transport for MockTransport {
        async fn read(&mut self) -> anyhow::Result<Option<Vec<u8>>> {
            if self.is_closed_flag {
                Ok(None)
            } else {
                Ok(self.read_data.take())
            }
        }
        
        async fn write(&mut self, data: &[u8]) -> anyhow::Result<()> {
            self.write_data.extend_from_slice(data);
            Ok(())
        }
        
        fn is_closed(&self) -> bool {
            self.is_closed_flag
        }

        async fn shutdown(&mut self) -> anyhow::Result<()> {
            self.is_closed_flag = true;
            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use mock::MockTransport;

    #[tokio::test]
    async fn test_tcp_transport_new() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        
        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let transport = TcpTransport::new(stream);
            assert!(!transport.is_closed());
        });

        let client_handle = tokio::spawn(async move {
            let stream = TcpStream::connect(addr).await.unwrap();
            let transport = TcpTransport::new(stream);
            assert!(!transport.is_closed());
        });

        // Use timeout to prevent hanging
        let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
        
        server_result.unwrap().unwrap();
        client_result.unwrap().unwrap();
    }

    #[tokio::test]
    async fn test_tcp_transport_write_read() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let mut transport = TcpTransport::new(stream);
            
            // Write some data
            Transport::write(&mut transport, b"test data").await.unwrap();
            
            // Read data back
            let data = Transport::read(&mut transport).await.unwrap();
            assert!(data.is_none()); // EOF since client closed
        });

        let client_handle = tokio::spawn(async move {
            let stream = TcpStream::connect(addr).await.unwrap();
            let mut transport = TcpTransport::new(stream);
            
            // Read data
            let data = Transport::read(&mut transport).await.unwrap().unwrap();
            assert_eq!(data, b"test data");
            
            // Close connection by dropping
        });

        // Use timeout to prevent hanging
        let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
        
        server_result.unwrap().unwrap();
        client_result.unwrap().unwrap();
    }

    #[tokio::test]
    async fn test_tcp_transport_closed() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let transport = TcpTransport::new(stream);
            
            // Close the connection from server side
            drop(transport);
        });

        let client_handle = tokio::spawn(async move {
            let stream = TcpStream::connect(addr).await.unwrap();
            let mut transport = TcpTransport::new(stream);
            
            // Try to read - should get None when connection is closed
            let data = Transport::read(&mut transport).await.unwrap();
            assert!(data.is_none());
            assert!(transport.is_closed());
        });

        // Use timeout to prevent hanging
        let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
        
        server_result.unwrap().unwrap();
        client_result.unwrap().unwrap();
    }

    #[tokio::test]
    async fn test_websocket_transport_new() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let ws_stream = accept_async(stream).await.unwrap();
            let transport = WebSocketTransport::new(ws_stream);
            assert!(!transport.is_closed());
        });

        let client_handle = tokio::spawn(async move {
            let (ws_stream, _) = connect_async(format!("ws://{}", addr)).await.unwrap();
            let transport = WebSocketTransport::new(ws_stream);
            assert!(!transport.is_closed());
        });

        // Use timeout to prevent hanging
        let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
        
        server_result.unwrap().unwrap();
        client_result.unwrap().unwrap();
    }

    #[tokio::test]
    async fn test_websocket_transport_send_metadata() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let ws_stream = accept_async(stream).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Read the metadata
            let data = Transport::read(&mut transport).await.unwrap().unwrap();
            let json: serde_json::Value = serde_json::from_slice(&data).unwrap();
            assert_eq!(json["host"], "example.com");
            assert_eq!(json["port"], 80);
        });

        let client_handle = tokio::spawn(async move {
            let (ws_stream, _) = connect_async(format!("ws://{}", addr)).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            transport.send_metadata("example.com", 80).await.unwrap();
        });

        // Use timeout to prevent hanging
        let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
        
        server_result.unwrap().unwrap();
        client_result.unwrap().unwrap();
    }

    #[tokio::test]
    async fn test_websocket_transport_write_read() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let ws_stream = accept_async(stream).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Read metadata first
            Transport::read(&mut transport).await.unwrap();
            
            // Read actual data
            let data = Transport::read(&mut transport).await.unwrap().unwrap();
            assert_eq!(data, b"websocket data");
            
            // Echo it back
            Transport::write(&mut transport, b"echo: websocket data").await.unwrap();
        });

        let client_handle = tokio::spawn(async move {
            let (ws_stream, _) = connect_async(format!("ws://{}", addr)).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Send metadata
            transport.send_metadata("example.com", 80).await.unwrap();
            
            // Send data
            Transport::write(&mut transport, b"websocket data").await.unwrap();
            
            // Read echo
            let data = Transport::read(&mut transport).await.unwrap().unwrap();
            assert_eq!(data, b"echo: websocket data");
        });

        // Use timeout to prevent hanging
        let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
        
        server_result.unwrap().unwrap();
        client_result.unwrap().unwrap();
    }

    #[tokio::test]
    async fn test_websocket_transport_closed() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let ws_stream = accept_async(stream).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Read metadata
            Transport::read(&mut transport).await.unwrap();
            
            // Close connection by shutting down transport
            transport.shutdown().await.unwrap();
        });

        let client_handle = tokio::spawn(async move {
            let (ws_stream, _) = connect_async(format!("ws://{}", addr)).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Send metadata
            transport.send_metadata("example.com", 80).await.unwrap();
            
            // Try to read - should eventually get None when connection is closed
            let data = Transport::read(&mut transport).await.unwrap();
            assert!(data.is_none());
            assert!(transport.is_closed());
        });

        // Use timeout to prevent hanging
        let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
        
        server_result.unwrap().unwrap();
        client_result.unwrap().unwrap();
    }

    #[tokio::test]
    async fn test_websocket_transport_text_message() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let ws_stream = accept_async(stream).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Read metadata
            Transport::read(&mut transport).await.unwrap();
            
            // Read text message
            let data = Transport::read(&mut transport).await.unwrap().unwrap();
            assert_eq!(String::from_utf8(data).unwrap(), "text message");
            transport.shutdown().await.unwrap();
        });

        let client_handle = tokio::spawn(async move {
            let (ws_stream, _) = connect_async(format!("ws://{}", addr)).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Send metadata
            transport.send_metadata("example.com", 80).await.unwrap();
            
            let text_data = "text message".as_bytes();
            Transport::write(&mut transport, text_data).await.unwrap();
            
            // The server should close the connection, so we should read None.
            let data = Transport::read(&mut transport).await.unwrap();
            assert!(data.is_none());
        });

        // Use timeout to prevent hanging
        let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
        
        server_result.unwrap().unwrap();
        client_result.unwrap().unwrap();
    }

    #[tokio::test]
    async fn test_websocket_transport_large_data() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        
        // Create large data payload
        let large_data = vec![0x42u8; 10000];
        let large_data_clone = large_data.clone();

        let server_handle = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let ws_stream = accept_async(stream).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Read metadata
            Transport::read(&mut transport).await.unwrap();
            
            // Read large data
            let data = Transport::read(&mut transport).await.unwrap().unwrap();
            assert_eq!(data.len(), 10000);
            assert_eq!(data, large_data_clone);
            
            // Echo it back
            Transport::write(&mut transport, &large_data_clone).await.unwrap();
        });

        let client_handle = tokio::spawn(async move {
            let (ws_stream, _) = connect_async(format!("ws://{}", addr)).await.unwrap();
            let mut transport = WebSocketTransport::new(ws_stream);
            
            // Send metadata
            transport.send_metadata("example.com", 80).await.unwrap();
            
            // Send large data
            Transport::write(&mut transport, &large_data).await.unwrap();
            
            // Read echo
            let data = Transport::read(&mut transport).await.unwrap().unwrap();
            assert_eq!(data.len(), 10000);
            assert_eq!(data, large_data);
        });

        // Use timeout to prevent hanging
        let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
        
        server_result.unwrap().unwrap();
        client_result.unwrap().unwrap();
    }
    
    #[tokio::test]
    async fn test_mock_transport_read_write() {
        let mut mock_transport = MockTransport::with_read_data(b"mock data".to_vec());
        
        // Test reading
        let data = Transport::read(&mut mock_transport).await.unwrap();
        assert!(data.is_some());
        assert_eq!(data.unwrap(), b"mock data");
        
        // Test writing
        Transport::write(&mut mock_transport, b"written data").await.unwrap();
        assert_eq!(mock_transport.write_data, b"written data");
    }
    
    #[tokio::test]
    async fn test_mock_transport_closed() {
        let mock_transport = MockTransport::closed();
        assert!(mock_transport.is_closed());
        
        let mut mock_transport_mut = MockTransport::closed();
        let data = Transport::read(&mut mock_transport_mut).await.unwrap();
        assert!(data.is_none());
    }
    
    #[tokio::test]
    async fn test_mock_transport_empty_read() {
        let mut mock_transport = MockTransport::new();
        let data = Transport::read(&mut mock_transport).await.unwrap();
        assert!(data.is_none());
    }
    
    // Test with different parameter combinations
    #[tokio::test]
    async fn test_transport_write_read_different_sizes() {
        let test_cases = vec![
            vec![], // Empty data
            vec![0x01], // Single byte
            vec![0x01, 0x02, 0x03], // Small data
            vec![0x42; 1024], // Medium data (1KB)
            vec![0x42; 8192], // Large data (8KB)
        ];
        
        for test_data in test_cases {
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();
            let test_data_clone = test_data.clone();
            
            let server_handle = tokio::spawn(async move {
                let (stream, _) = listener.accept().await.unwrap();
                let mut transport = TcpTransport::new(stream);
                
                // Write test data
                if !test_data_clone.is_empty() {
                    Transport::write(&mut transport, &test_data_clone).await.unwrap();
                }
                // Shutdown the write side of the stream to signal EOF to the client.
                transport.shutdown().await.unwrap();
            });
            
            let client_handle = tokio::spawn(async move {
                let stream = TcpStream::connect(addr).await.unwrap();
                let mut transport = TcpTransport::new(stream);
                
                // Read test data
                let mut received_data = Vec::new();
                while let Some(chunk) = Transport::read(&mut transport).await.unwrap() {
                    received_data.extend_from_slice(&chunk);
                }
                assert_eq!(received_data, test_data);
            });
            
            // Use timeout to prevent hanging
            let server_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
            let client_result = tokio::time::timeout(tokio::time::Duration::from_secs(5), client_handle).await;
            
            server_result.unwrap().unwrap();
            client_result.unwrap().unwrap();
        }
    }
}